import rclpy
from rclpy.node import Node
import numpy as np
import open3d as o3d
from sensor_msgs.msg import PointCloud2
from sensor_msgs_py import point_cloud2
import math
import time

class CableStatusAnalyzer(Node):
    def __init__(self):
        super().__init__('cable_status_analyzer')
        self.subscription = self.create_subscription(
            PointCloud2,
            'simulated_pointcloud',
            self.listener_callback,
            10)
        self.subscription
        
        # 可视化配置
        self.vis = o3d.visualization.VisualizerWithKeyCallback()
        self.vis.create_window(window_name="Cable Analyzer")
        opt = self.vis.get_render_option()
        opt.background_color = [0, 0, 0]
        opt.point_size = 3
        
        # 数据存储
        self.clustered_pcd = o3d.geometry.PointCloud()
        self.coord_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
        self.vis.add_geometry(self.coord_frame)
        
        # 分析结果
        self.cable_params = {
            'left': {'length': 0.0, 'angle': 0.0, 'diameter': 0.0},
            'right': {'length': 0.0, 'angle': 0.0, 'diameter': 0.0}
        }
        self.cargo_pos = [0.0, 0.0, 0.0]
        self.first_update = True
        
        # 处理时间跟踪
        self.processing_time_ms = 0.0  # 新增：存储处理时间（毫秒）

        # 聚类关键参数（针对全方位摆动优化）
        self.dbscan_eps = 0.3  # 密度聚类的邻域半径（根据钢缆直径调整）
        self.dbscan_min_samples = 15  # 最小点数（控制聚类粒度）
        self.min_cable_points = 300  # 钢缆最小点数阈值

    def split_by_z(self, points, z_threshold_ratio=0.85):
        """分割货物（下方）和钢缆（上方）"""
        z_values = points[:, 2]
        z_max = np.max(z_values)
        cargo_mask = z_values > z_max * z_threshold_ratio
        return points[cargo_mask], points[~cargo_mask]

    def cluster_cables_by_density(self, cable_points):
        """基于密度的钢缆聚类（解决交叉摆动问题）
        原理：钢缆是连续的高密度点集，即使摆动交叉，单根钢缆的点仍保持局部连续性
        """
        if len(cable_points) < self.min_cable_points * 2:
            self.get_logger().warn(f"钢缆点不足({len(cable_points)})")
            return None, None
        
        # 1. 转换为Open3D点云
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(cable_points)
        
        # 2. 密度聚类（DBSCAN）- 自动识别连续的点集
        with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Error):
            labels = np.array(pcd.cluster_dbscan(
                eps=self.dbscan_eps,
                min_points=self.dbscan_min_samples,
                print_progress=False
            ))
        
        # 3. 提取两个最大的聚类（对应两根钢缆）
        unique_labels, counts = np.unique(labels, return_counts=True)
        # 排除噪声点（label=-1）
        valid_labels = unique_labels[unique_labels != -1]
        if len(valid_labels) < 2:
            self.get_logger().warn(f"仅识别到{len(valid_labels)}根钢缆，使用备用方案")
            return self.fallback_split(cable_points)
        
        # 按点数排序，取前两个最大聚类
        sorted_indices = np.argsort(counts[unique_labels != -1])[::-1]
        main_labels = valid_labels[sorted_indices[:2]]
        
        cluster1 = cable_points[labels == main_labels[0]]
        cluster2 = cable_points[labels == main_labels[1]]
        
        # 4. 标记左右（基于上部点的X坐标，不受摆动影响）
        return self.label_left_right(cluster1, cluster2)

    def fallback_split(self, cable_points):
        """备用分割方案：当密度聚类失败时使用"""
        # 取上部30%的点（靠近雷达，摆动影响小）
        z_values = cable_points[:, 2]
        z_min = np.min(z_values)
        z_max = np.max(z_values)
        top_mask = z_values < z_min + (z_max - z_min) * 0.3
        top_points = cable_points[top_mask]
        
        if len(top_points) < 100:
            x_median = np.median(cable_points[:, 0])
            return cable_points[cable_points[:,0]<x_median], cable_points[cable_points[:,0]>=x_median]
        
        # 基于上部点的X中心分割
        x_median_top = np.median(top_points[:, 0])
        return cable_points[cable_points[:,0]<x_median_top], cable_points[cable_points[:,0]>=x_median_top]

    def label_left_right(self, cluster1, cluster2):
        """标记左右钢缆（基于上部稳定点）"""
        # 取每个聚类上部20%的点（摆动影响小）
        def get_top_x(cluster):
            z_values = cluster[:, 2]
            z_min = np.min(z_values)
            z_max = np.max(z_values)
            top_mask = z_values < z_min + (z_max - z_min) * 0.2
            return np.median(cluster[top_mask, 0])  # 上部点的X中位数
        
        x1 = get_top_x(cluster1)
        x2 = get_top_x(cluster2)
        
        # X小的为左钢缆，X大的为右钢缆
        if x1 < x2:
            return cluster1, cluster2
        else:
            return cluster2, cluster1

    def perform_clustering(self, points):
        """主聚类流程：先分货物和钢缆，再用密度聚类分钢缆"""
        cargo_points, cable_points = self.split_by_z(points)
        
        if len(cable_points) < self.min_cable_points * 2:
            return None, None, None
        
        left_cable, right_cable = self.cluster_cables_by_density(cable_points)
        
        # 过滤过小聚类
        if len(left_cable) < self.min_cable_points:
            left_cable = np.array([])
        if len(right_cable) < self.min_cable_points:
            right_cable = np.array([])
        
        self.get_logger().info(
            f"聚类完成: 左钢缆={len(left_cable)}, 右钢缆={len(right_cable)}, 货物={len(cargo_points)}"
        )
        return left_cable, right_cable, cargo_points

    def fit_line(self, points):
        """拟合钢缆直线"""
        if len(points) < 50:
            return None, None
        mean = np.mean(points, axis=0)
        centered = points - mean
        u, s, vh = np.linalg.svd(centered)
        direction = vh[0]
        return direction, mean

    def analyze_cable(self, cable_points, is_left=True):
        """分析钢缆参数"""
        if len(cable_points) < self.min_cable_points:
            return {'length': 0.0, 'angle': 0.0, 'diameter': 0.0}
        
        direction, mean = self.fit_line(cable_points)
        if direction is None:
            return {'length': 0.0, 'angle': 0.0, 'diameter': 0.0}
        
        # 修正方向向量（向下）
        if direction[2] < 0:
            direction = -direction
        
        # 长度计算
        projections = np.dot(cable_points - mean, direction)
        length = np.max(projections) - np.min(projections)
        
        # 偏角计算
        vertical_dir = np.array([0, 0, 1])
        dot_product = np.dot(direction, vertical_dir)
        dot_product = np.clip(dot_product, -1.0, 1.0)
        angle_deg = np.degrees(np.arccos(dot_product))
        
        # 直径计算
        centered = cable_points - mean
        perpendicular = centered - np.outer(projections, direction)
        distances = np.linalg.norm(perpendicular, axis=1)
        diameter = 2 * np.mean(distances)
        
        return {'length': length, 'angle': angle_deg, 'diameter': diameter}

    def visualize_results(self, left, right, cargo):
        """可视化聚类结果（左红右绿）"""
        all_points = []
        all_colors = []
        
        if len(left) > 0:
            all_points.append(left)
            all_colors.append(np.tile([1, 0, 0], (len(left), 1)))  # 红色
        
        if len(right) > 0:
            all_points.append(right)
            all_colors.append(np.tile([0, 1, 0], (len(right), 1)))  # 绿色
        
        if len(cargo) > 0:
            all_points.append(cargo)
            all_colors.append(np.tile([0, 0, 1], (len(cargo), 1)))  # 蓝色
        
        if not all_points:
            return
        
        all_points = np.vstack(all_points)
        all_colors = np.vstack(all_colors)
        
        self.clustered_pcd.points = o3d.utility.Vector3dVector(all_points)
        self.clustered_pcd.colors = o3d.utility.Vector3dVector(all_colors)
        
        if self.first_update:
            self.vis.add_geometry(self.clustered_pcd)
            self.first_update = False
        else:
            self.vis.update_geometry(self.clustered_pcd)
        
        self.vis.poll_events()
        self.vis.update_renderer()

    def display_analysis(self):
        """显示分析结果和处理时间"""
        self.get_logger().info("\n===== 钢缆状态分析 =====")
        self.get_logger().info(
            f"左钢缆: 长度={self.cable_params['left']['length']:.2f}m, "
            f"偏角={self.cable_params['left']['angle']:.2f}°, "
            f"直径={self.cable_params['left']['diameter']:.4f}m"
        )
        self.get_logger().info(
            f"右钢缆: 长度={self.cable_params['right']['length']:.2f}m, "
            f"偏角={self.cable_params['right']['angle']:.2f}°, "
            f"直径={self.cable_params['right']['diameter']:.4f}m"
        )
        self.get_logger().info(f"处理时间: {self.processing_time_ms:.2f} ms")  # 修改：显示处理时间
        self.get_logger().info("=======================\n")

    def listener_callback(self, msg):
        """主回调函数"""
        # 记录开始时间
        start_time = time.time()
        
        points = point_cloud2.read_points_list(msg, field_names=("x", "y", "z"), skip_nans=True)
        point_array = np.array([(p.x, p.y, p.z) for p in points], dtype=np.float32)
        self.get_logger().info(f"接收点云: {len(point_array)}个点")
        
        if len(point_array) < 1000:
            self.get_logger().warn("点云数量不足，跳过分析")
            return
        
        left_cable, right_cable, cargo_points = self.perform_clustering(point_array)
        if left_cable is None:
            return
        
        self.cable_params['left'] = self.analyze_cable(left_cable)
        self.cable_params['right'] = self.analyze_cable(right_cable)
        self.cargo_pos = np.mean(cargo_points, axis=0) if len(cargo_points) > 0 else [0,0,0]
        
        self.visualize_results(left_cable, right_cable, cargo_points)
        
        # 计算并存储处理时间（毫秒）
        end_time = time.time()
        self.processing_time_ms = (end_time - start_time) * 1000  # 转换为毫秒
        
        self.display_analysis()

def main(args=None):
    rclpy.init(args=args)
    analyzer = CableStatusAnalyzer()
    try:
        rclpy.spin(analyzer)
    except KeyboardInterrupt:
        analyzer.vis.destroy_window()
    analyzer.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
